import json

from flask import Flask, request, jsonify, render_template

from news_classification.glda import get_glda_infer_engine, infer_doc_topics_by_glda, read_topic_words

app = Flask(__name__)

model_fp = 'model/policy/256-topics-lda-1k-iterations.model'
topic_num = 256
topics_assign = {
    # 9: '抽查检查', 92: '产品质量', 148: '抽查检查', 227: '抽查',
    # 189: '违法执法', 60: '专项整治', 131: '行政执法', 13: '违法处罚',
    114: '中小学教育', 204: '教育培训', 61: '教育', 248: '体育', 46: '招生', 48: '科普',  # 文体教育
    # 167: '科技创新',
    51: '水利',
    # 73: '生产安全', 171: '事故应急',
    # 209: '应急救援', 157: '自然灾害', 12: '地质灾害',
    134: '工程建设', 47: '工程建设',
    95: '房产建筑', 76: '用地规划', 83: '土地用地', 232: '房地产', 133: '用地', 147: '征地拆迁',
    251: '环评', 156: '环境保护', 4: '环境整治', 27: '污染排放', 96: '节能减排', 66: '污水处理', 119: '生态保护', 161: '垃圾处理', 196: '污染治理',
    222: '交通运输', 25: '道路交通', 172: '公共交通', 85: '交通运输',
    # 128: '资金补助',
    # 219: '收费定价',
    # 145: '招投标', 59: '采购',
    175: '农业', 212: '渔业', 98: '农产品', 192: '林业', 193: '农业机械', 231: '粮油',  # 农林牧渔
    39: '医疗', 152: '医院', 26: '健康', 33: '医疗医保', 229: '传染病',
    30: '食品药品', 127: '医药',
    # 198: '招商引资',
    # 49: '考核考评',
    # 77: '评选奖励',
    205: '旅游', 11: '文化旅游',
    159: '石油',
    115: '金融', 151: '银行', 149: '基金投资',
    # 228: '税务',
    207: '疫情防控',
    63: '气象',
    # 253: '社保',
    200: '电力',
    # 186: '信息化建设', 99: '信息化建设',
    68: '文化产业', 249: '广电', 70: '历史文物',
    # 178: '燃气供水',
    71: '国资国企',
    176: '就业',
    # 94: '救助',
    # 75: '农村建设',
    146: '养老',
    153: '扶贫',
    22: '进出口',
    # 162: '出入境',
    # 15: '残疾人',
    182: '诚信信用',
    218: '矿产资源', 247: '煤矿',
    44: '危化品',
}

categories = {
    '文体教育': {114: '中小学教育', 204: '教育培训', 61: '教育', 248: '体育', 46: '招生', 48: '科普'},
    '水利': {51: '水利'},
    '工程建设': {134: '工程建设', 47: '工程建设'},
    '房产建筑': {95: '房产建筑', 76: '用地规划', 83: '土地用地', 232: '房地产', 133: '用地', 147: '征地拆迁'},
    '环保行业': {251: '环评', 156: '环境保护', 4: '环境整治', 27: '污染排放', 96: '节能减排', 66: '污水处理',
             119: '生态保护', 161: '垃圾处理', 196: '污染治理'},
    '交通运输': {222: '交通运输', 25: '道路交通', 172: '公共交通', 85: '交通运输'},
    '农林牧渔': {175: '农业', 212: '渔业', 98: '农产品', 192: '林业', 193: '农业机械', 231: '粮油'},
    '医疗健康': {39: '医疗', 152: '医院', 26: '健康', 33: '医疗医保', 229: '传染病'},
    '食品药品': {30: '食品药品', 127: '医药'},
    '旅游': {205: '旅游', 11: '文化旅游'},
    '石油': {159: '石油'},
    '金融': {115: '金融', 151: '银行', 149: '基金投资'},
    '气象': {63: '气象'},
    '电力': {200: '电力'},
    '文化产业': {68: '文化产业', 249: '广电', 70: '历史文物'},
    '进出口': {22: '进出口'},
    '矿产': {218: '矿产资源', 247: '煤矿'},
    '危化品': {44: '危化品'},
    '国资国企': {71: '国资国企'},
    '就业': {176: '就业'},
    '养老': {146: '养老'},
    '诚信信用': {182: '诚信信用'},
}
inf_engine = get_glda_infer_engine(model_fp)
topic_words = read_topic_words(model_fp)


def get_assign_topic(topic_id):
    try:
        return topics_assign[topic_id]
    except KeyError:
        return 'Topic %d' % topic_id


def get_assign_category(topic_id):
    for category, topics in categories.items():
        if topic_id in topics:
            return category
    return '其他'


@app.route('/infer_doc_topics', methods=['POST'])
def infer_doc_topics():
    doc = request.form['doc']
    topic_dis = infer_doc_topics_by_glda(doc, inf_engine)
    result = []
    for topic_id, prob in topic_dis:
        if prob < 0.1:
            continue
        result.append(
            {
                'topic_id': topic_id,
                'topic_name': get_assign_topic(topic_id),
                'topic_weight': prob,
                'topic_words': ', '.join([e[0] for e in topic_words[topic_id]]),
                'category': get_assign_category(topic_id)
            }
        )
    return jsonify(result)


@app.route('/infer_doc_categories', methods=['POST'])
def infer_doc_categories():
    doc = request.form['doc']
    topic_dis = infer_doc_topics_by_glda(doc, inf_engine)
    result = []
    for topic_id, prob in topic_dis:
        if prob < 0.1:
            continue
        result.append(
            {
                'topic_id': topic_id,
                'topic_name': get_assign_topic(topic_id),
                'topic_weight': prob,
                'topic_words': ', '.join([e[0] for e in topic_words[topic_id]]),
                'category': get_assign_category(topic_id)
            }
        )
    if result:
        categories = list(set([r['category'] for r in result if r['category'] != '其他']))
        if len(categories) == 0:
            categories = ['其他']
    else:
        categories = ['其他']
    return jsonify({'categories': categories})


@app.route('/get_topic_words', methods=['GET'])
def get_topic_words():
    return jsonify(topic_words)


def read_samples():
    sample_policy_content = []
    with open('data/policy/policy-content-1w-sample.txt', encoding='utf-8') as fi:
        for line in fi:
            sample_policy_content.append(line.strip())
    sample_topics = {}
    with open('data/policy/policy-1w-sample-topics.jsonl', encoding='utf-8') as fi:
        for i, line in enumerate(fi):
            row_res = json.loads(line)
            for e in row_res:
                if e['topic_id'] not in sample_topics:
                    sample_topics[e['topic_id']] = [(i, e['topic_weight'])]
                else:
                    sample_topics[e['topic_id']].append((i, e['topic_weight']))
    sample_categories = {}
    with open('data/policy/policy-content-1w-sample-category.txt', encoding='utf-8') as fi:
        for i, line in enumerate(fi):
            row_categories = line.strip().split(', ')
            for category in row_categories:
                if category not in sample_categories:
                    sample_categories[category] = [i]
                else:
                    sample_categories[category].append(i)
    return sample_policy_content, sample_topics, sample_categories


sample_policy_content, sample_topics, sample_categories = read_samples()


@app.route('/show_sample_topic_policy', methods=['GET'])
def show_sample_topic_policy():
    topic_id = request.args.get('topic_id', type=int)
    sample_topic_words = ', '.join([e[0] for e in topic_words[topic_id]])
    marked_policy = [(sample_policy_content[i], weight) for i, weight
                     in sorted(sample_topics[topic_id], key=lambda x: x[1], reverse=True)]
    return render_template(
        'show_sample_topic_policy.html',
        topic_id=topic_id,
        sample_topic_words=sample_topic_words,
        marked_policy=marked_policy)


@app.route('/show_sample_topics', methods=['GET'])
def show_sample_topics():
    data = []
    for topic_id, sample_ids in sample_topics.items():
        sample_topic_words = ', '.join([e[0] for e in topic_words[topic_id]])
        row = [topic_id, sample_topic_words, len(sample_ids)]
        data.append(row)
    data = sorted(data, key=lambda x: x[2], reverse=True)
    return render_template('show_sample_topics.html', data=data)


@app.route('/show_sample_category_policy', methods=['GET'])
def show_sample_category_policy():
    category = request.args.get('category', type=str)
    marked_policy = [sample_policy_content[i] for i in sample_categories[category]]
    return render_template(
        'show_sample_category_policy.html',
        category=category,
        marked_policy=marked_policy)


@app.route('/show_sample_categories', methods=['GET'])
def show_sample_categories():
    data = []
    for category, sample_ids in sample_categories.items():
        row = [category, len(sample_ids)]
        data.append(row)
    data = sorted(data, key=lambda x: x[1], reverse=True)
    return render_template('show_sample_categories.html', data=data)


if __name__ == '__main__':
    app.run()
